package cn.flying.cloud.base.core.filter.wrapper;

import javax.servlet.ReadListener;
import javax.servlet.ServletException;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.Part;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Collection;
import java.util.Locale;
import java.util.Map;

import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.text.StringEscapeUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StreamUtils;

/**
 * 自定义http请求对象，解决请求body不能多次读取的问题
 * 可以参照spring boot 内置的ContentCachingRequestWrapper
 *
 * @author: admin
 * @date: 2023年05月28日 11:23
 * @version: 1.0
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
    private final Logger logger = LoggerFactory.getLogger(this.getClass());

    private byte[] requestBody;
    private Map<String, String[]> params;
    private Collection<Part> parts;

    private static final String POST_METHOD = "POST";
    private static final String MULTIPART = "multipart/";

    public XssHttpServletRequestWrapper(HttpServletRequest request) throws IOException, ServletException {
        super(request);
        this.params = request.getParameterMap();
        if (isMultipartContent(request)) {
            this.parts = request.getParts();
        } else {
            this.requestBody = StreamUtils.copyToByteArray(request.getInputStream());
        }
    }

    private boolean isMultipartContent(HttpServletRequest request) {
        if (!POST_METHOD.equalsIgnoreCase(request.getMethod())) {
            return false;
        }
        String contentType = request.getContentType();
        if (contentType == null) {
            return false;
        }
        if (contentType.toLowerCase(Locale.ENGLISH).startsWith(MULTIPART)) {
            return true;
        }
        return false;
    }

    /**
     * xss攻击防范处理
     *
     * @param name
     */
    @Override
    public String getParameter(String name) {
        logger.info("getParameter方法xss攻击防范处理");
        String[] values = params.get(name);
        if (ArrayUtils.isEmpty(values)) {
            return null;
        }
        return values[0];
    }

    /**
     * xss攻击防范处理
     *
     * @param name
     */
    @Override
    public String[] getParameterValues(String name) {
        logger.info("getParameterValues方法xss攻击防范处理");
        String[] values = params.get(name);
        if (values != null && values.length > 0) {
            for (int i = 0; i < values.length; i++) {
                values[i] = StringEscapeUtils.escapeHtml4(values[i]);
            }
        }
        return values;
    }

    /**
     * xss攻击防范处理
     */
    @Override
    public Map<String, String[]> getParameterMap() {
        logger.info("getParameterMap方法xss攻击防范处理");
        Map<String, String[]> map = params;
        if (map != null && !map.isEmpty()) {
            for (String[] value : map.values()) {
                //循环所有的value
                for (int i = 0; i < value.length; i++) {
                    value[i] = StringEscapeUtils.escapeHtml4(value[i]);
                }
            }
        }
        return map;
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        final ByteArrayInputStream bis = new ByteArrayInputStream(requestBody);
        return new ServletInputStream() {
            private boolean finished = false;

            @Override
            public int read() throws IOException {
                return bis.read();
            }

            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public int available() throws IOException {
                return bis.available();
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setReadListener(ReadListener readListener) {

            }
        };
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }
}